import pygame as pg
import random

size = 20  # 格子大小
long = 5  # 蛇单位长
food_v = 1*size
White = (255, 255, 255)
Black = (0, 0, 0)
Blue = (0, 0, 255)
Red = (255, 0, 0)
Green = (0, 255, 0)
Gray = (128, 128, 128)


class Button:
    def __init__(self, x, y, text, color=Black):
        self.x = x*size
        self.y = y*size
        self.width = 200
        self.height = 100
        self.color = color
        self.font = pg.font.Font(None, 36)
        self.text = text

    def draw(self, screen):
        pg.draw.rect(screen, self.color, (self.x, self.y, self.width, self.height))  # 设置矩形的颜色和位置
        text_surface = self.font.render(self.text, True, White)  # 设置字体和颜色
        text_rect = text_surface.get_rect(center=(self.x+self.width/2, self.y+self.height/2))  # 设置文字位置
        screen.blit(text_surface, text_rect)


class Snake:
    def __init__(self, x, y, color, direction):
        self.x = x*size
        self.y = y*size
        self.direction = direction
        self.body = []
        self.width = size
        self.head_color = Red
        self.body_color = color
        self.length = long
        self.score = 0
        self.nex = True
        self.over = False
        for i in range(self.length):
            self.add()

    def add(self):
        node_x, node_y = (self.x, self.y)
        if self.body:
            node_x, node_y = self.body[0].x, self.body[0].y
        node = pg.Rect(node_x, node_y, self.width, self.width)
        if self.direction == pg.K_RIGHT:
            node.x = (node.x + size) % 800
        elif self.direction == pg.K_LEFT:
            node.x = (node.x - size) % 800
        elif self.direction == pg.K_UP:
            node.y = (node.y - size) % 600
        elif self.direction == pg.K_DOWN:
            node.y = (node.y + size) % 600
        self.body.insert(0, node)

    def delete(self):
        self.body.pop()  # 不加入参数自动移除末尾元素

    def increase(self):
        self.nex = False
        return True

    def move(self):
        self.add()
        if self.nex:
            self.delete()
        self.nex = True

    def is_dead(self):
        if self.score >= 100:
            self.over = True
            return "win"
        if self.body[0].x < 0 or self.body[0].x > 800 or self.body[0].y < 0 or self.body[0].y > 600:
            self.over = True
            return "dead"
        if self.body[0] in self.body[1:]:
            self.over = True
            return "dead"

    def draw(self, screen):
        if not self.is_dead() and not self.over:
            self.move()
        for nodes in self.body:
            pg.draw.rect(screen, self.body_color, nodes)
            pg.draw.rect(screen, self.head_color, self.body[0])


class Food:
    def __init__(self):
        self.x = random.randint(0, 39)*size
        self.y = random.randint(0, 29)*size
        self.all = []
        self.color = Black
        self.width = size
        self.nex = True
        self.all.insert(0, (self.x, self.y))

    def set(self):
        self.x = random.randint(0, 39)*size
        self.y = random.randint(0, 29)*size
        node = (self.x, self.y)
        if self.die():
            self.all.insert(0, node)

    def move(self):
        direct = random.randint(0, 3)
        num = food_v
        if self.nex:
            if direct == 0:
                self.x = (self.x + num) % 800
            elif direct == 1:
                self.x = (self.x - num) % 800
            elif direct == 2:
                self.y = (self.y - num) % 600
            elif direct == 3:
                self.y = (self.y + num) % 600

    def die(self):
        del self.all[0]
        return True

    def draw(self, screen):
        pg.draw.rect(screen, self.color, (self.x, self.y, self.width, self.width))


def init():
    pg.init()
    screen_width = 800
    screen_height = 600
    screen = pg.display.set_mode((screen_width, screen_height))
    pg.display.set_caption('secreted snake')
    button1 = Button(14, 9, 'start')
    button2 = Button(14, 15, 'exit')
    button3 = Button(14, 6, 'f')
    button4 = Button(14, 12, 'd')
    button5 = Button(14, 18, 'return')
    snake1 = Snake(9, 10, Blue, pg.K_RIGHT)
    snake2 = Snake(30, 20, Green, pg.K_LEFT)
    score1 = Button(5, 24, 'Score1: ' + str(snake1.score), Gray)
    score2 = Button(15, 24, 'Score2: ' + str(snake2.score), Gray)
    food1 = Food()
    clock = pg.time.Clock()
    running = True
    nexted = True  # 进入下一个界面
    playing = False  # 游戏开始
    game_type = True  # 选择游戏类型
    over = False  # 游戏结束
    while running:
        for event in pg.event.get():
            if event.type == pg.QUIT:  # 如果关闭窗口
                running = False
            elif event.type == pg.MOUSEBUTTONDOWN:  # 鼠标按下
                mouse_x, mouse_y = event.pos  # 获取鼠标位置
                if (button1.x < mouse_x < button1.x+button1.width and button1.y < mouse_y < button1.y+button1.height and
                        nexted and not playing):
                    nexted = False
                elif (button2.x < mouse_x < button2.x+button2.width and button2.y < mouse_y < button2.y+button2.height
                      and nexted and not playing):
                    running = False
                elif (button3.x < mouse_x < button3.x+button3.width and button3.y < mouse_y < button3.y+button3.height
                      and not nexted and not playing):
                    game_type = True
                    playing = True
                elif (button4.x < mouse_x < button4.x+button4.width and button4.y < mouse_y < button4.y+button4.height
                      and not nexted and not playing):
                    game_type = False
                    playing = True
                elif (button5.x < mouse_x < button5.x+button5.width and button5.y < mouse_y < button5.y+button5.height
                      and not nexted and not playing):
                    nexted = True
            elif event.type == pg.KEYDOWN and playing:
                if event.key == pg.K_RIGHT or event.key == pg.K_LEFT or event.key == pg.K_UP or event.key == pg.K_DOWN:
                    snake1.direction = event.key
                if not game_type:
                    if event.key == pg.K_a or event.key == pg.K_d or event.key == pg.K_w or event.key == pg.K_s:
                        change_key(event, snake2)
                if event.key == pg.K_SPACE and over:
                    return main()
        screen.fill(pg.Color('white'))
        if nexted:
            button1.draw(screen)
            button2.draw(screen)
        elif not nexted and not playing:
            button3.draw(screen)
            button4.draw(screen)
            button5.draw(screen)
        if playing:
            food1.draw(screen)
            food1.move()
            snake1.draw(screen)
            score1.draw(screen)
            eat_food(snake1, score1, food1)
            if snake1.is_dead():
                over = is_dead(screen, snake1)
                food1.nex = False
            if not game_type:
                snake2.draw(screen)
                score2.draw(screen)
                eat_food(snake2, score2, food1)
                if snake2.is_dead():
                    over = is_dead(screen, snake2)
                    food1.nex = False
                if not battle(snake1, snake2):
                    over = True
                    lose = Button(10, 10, 'You Lose', Gray)
                    lose.draw(screen)
                    snake1.over = True
                    snake2.over = True
                    food1.nex = False
        pg.display.update()
        clock.tick(20)
    pg.quit()


def is_dead(screen, snake):
    win = Button(10, 10, 'You Win', Gray)
    lose = Button(10, 10, 'You Lose', Gray)
    if snake.is_dead() == "win":
        win.draw(screen)
    elif snake.is_dead() == "dead":
        lose.draw(screen)
    return True


def battle(snake1, snake2):
    for i in range(0, len(snake1.body)):
        if snake2.body[0].colliderect(snake1.body[i]):
            return False
    for i in range(0, len(snake2.body)):
        if snake1.body[0].colliderect(snake2.body[i]):
            return False
    return True


def change_key(event, snake):
    if event.key == pg.K_a:
        snake.direction = pg.K_LEFT
    elif event.key == pg.K_d:
        snake.direction = pg.K_RIGHT
    elif event.key == pg.K_w:
        snake.direction = pg.K_UP
    elif event.key == pg.K_s:
        snake.direction = pg.K_DOWN


def eat_food(snake, score, food):
    if snake.body[0].x == food.x and snake.body[0].y == food.y:
        snake.score += 10
        snake.increase()
        score.text = 'Score: ' + str(snake.score)
        food.set()


def main():
    # 初始化
    init()


if __name__ == '__main__':
    main()
